//
// Created by 12992 on 24-10-26.
//
#include "rsa.h"

#include <iostream>
#include <cstdio>
#include <openssl/err.h>
#include <openssl/pem.h>

#if _WIN32
extern "C"
{
#include "openssl/applink.c"
};
#endif


namespace crypto
{
    using namespace std;

    void rsa::generate_key_pair(const std::filesystem::path& private_key,
                                const std::filesystem::path& public_key)
    {
        // 初始化openssl加解密库
        OPENSSL_init_crypto(OPENSSL_INIT_LOAD_CONFIG, nullptr);

        // 创建密钥对上下文
        EVP_PKEY_CTX* ctx = EVP_PKEY_CTX_new_id(EVP_PKEY_RSA, nullptr);
        if (!ctx)
        {
            cerr << "can not create evp ctx" << endl;
            return;
        }

        // 初始化密钥对上下文
        int ret = EVP_PKEY_keygen_init(ctx);
        if (ret <= 0)
        {
            cerr << "can not init evp ctx" << endl;
            goto CLEAR_CTX;
        }

        // 设置密钥对长度
        ret = EVP_PKEY_CTX_set_rsa_keygen_bits(ctx, 2048);
        if (ret <= 0)
        {
            cerr << "can not set rsa keygen bit" << endl;
            goto CLEAR_CTX;
        }

        // 生成密钥对
        EVP_PKEY* key_pair{};
        ret = EVP_PKEY_keygen(ctx, &key_pair);
        if (ret <= 0)
        {
            cerr << "can not generate key pair" << endl;
            goto CLEAR_CTX;
        }

        // 打开私钥文件
        FILE* private_key_file{};
        ret = fopen_s(&private_key_file, private_key.string().c_str(), "wb");
        if (ret)
        {
            cerr << "can not open private key file" << endl;
            goto CLEAR_KEY;
        }

        // 打开公钥文件
        FILE* public_key_file{};
        ret = fopen_s(&public_key_file, public_key.string().c_str(), "wb");
        if (ret)
        {
            cerr << "can not open public key file" << endl;
            goto CLOSE_PRIVATE;
        }

        //写入私钥
        ret = PEM_write_PrivateKey(private_key_file, key_pair, nullptr, nullptr, 0, nullptr, nullptr);
        if (ret <= 0)
        {
            cerr << "can not write private key" << endl;
            goto CLOSE_PUBLIC;
        }

        // 写入公钥
        ret = PEM_write_PUBKEY(public_key_file, key_pair);
        if (ret <= 0)
        {
            cerr << "can not write public key" << endl;
            goto CLOSE_PUBLIC;
        }

        cout << "RSA key pair successfully generated" << endl;

    CLOSE_PUBLIC:
        fclose(public_key_file);
    CLOSE_PRIVATE:
        fclose(private_key_file);
    CLEAR_KEY:
        EVP_PKEY_free(key_pair);
    CLEAR_CTX:
        EVP_PKEY_CTX_free(ctx);
    }

    rsa::rsa(key_type type, const std::filesystem::path& key)
    {
        // 初始化openssl加解密库

        enum class error_type
        {
            NONE,
            OPENSSL_INIT_ERROR,
            KEY_OPEN_ERROR,
            KEY_READ_ERROR,
            CTX_CREATE_ERROR,
        };
        auto err = error_type::NONE;
        do
        {
            FILE* key_fd{};
            // 初始化openssl
            int ret = OPENSSL_init_crypto(OPENSSL_INIT_LOAD_CONFIG, nullptr);
            if (ret <= 0)
            {
                err = error_type::OPENSSL_INIT_ERROR;
                break;
            }

            // 打开密钥文件
            ret = fopen_s(&key_fd, key.string().c_str(), "rb");
            if (ret < 0)
            {
                err = error_type::KEY_OPEN_ERROR;
                break;
            }

            // 读取密钥
            if (key_type::PUBLIC == type)
            {
                // 公钥
                this->key = PEM_read_PUBKEY(key_fd, nullptr, nullptr, nullptr);
            }
            else
            {
                // 私钥
                this->key = PEM_read_PrivateKey(key_fd, nullptr, nullptr, nullptr);
            }
            fclose(key_fd);
            if (!this->key)
            {
                err = error_type::KEY_READ_ERROR;
                break;
            }

            // 创建加解密上下文
            this->ctx = EVP_PKEY_CTX_new(this->key, nullptr);
            if (!ctx)
            {
                err = error_type::CTX_CREATE_ERROR;
                break;
            }
        }
        while (false);
        switch (err)
        {
        case error_type::CTX_CREATE_ERROR:
            EVP_PKEY_free(this->key);
        case error_type::KEY_READ_ERROR: break;
        case error_type::KEY_OPEN_ERROR: break;
        case error_type::OPENSSL_INIT_ERROR: break;
        case error_type::NONE: ;
        }
    }

    rsa::~rsa()
    {
        if (this->key)
        {
            EVP_PKEY_free(this->key);
            EVP_PKEY_CTX_free(this->ctx);
        }
    }

    rsa::rsa(rsa&& obj) noexcept
        : key(exchange(obj.key, nullptr)), ctx(exchange(obj.ctx, nullptr))
    {
    }

    rsa& rsa::operator=(rsa&& obj) noexcept
    {
        this->key = exchange(obj.key, nullptr);
        this->ctx = exchange(obj.ctx, nullptr);
        return *this;
    }

    std::string rsa::encrypt(const std::string& data)
    {
        // 初始化加密操作
        int ret = EVP_PKEY_encrypt_init(this->ctx);
        if (ret <= 0)
        {
            throw runtime_error("can not init encrypt");
        }

        // 获取密文大小
        size_t out_size{};
        auto data_ptr = reinterpret_cast<const unsigned char*>(data.c_str());
        ret = EVP_PKEY_encrypt(this->ctx, nullptr, &out_size, data_ptr, data.length());
        if (ret <= 0)
        {
            throw runtime_error("can not get encrypt data size");
        }

        // 加密
        string output(out_size, '\0');
        const auto out_ptr = reinterpret_cast<unsigned char*>(output.data());
        ret = EVP_PKEY_encrypt(this->ctx, out_ptr, &out_size, data_ptr, data.length());
        if (ret <= 0)
        {
            // 获取错误信息
            unsigned long err_code = 0;
            while ((err_code = ERR_get_error()))
            {
                char* err_str = ERR_error_string(err_code, nullptr);
                fprintf(stderr, "Decryption error: %s\n", err_str);
            }
            throw runtime_error("can not decrypt data");
        }
        return output;
    }

    std::string rsa::decrypt(const std::string& data)
    {
        // 初始化解密操作
        int ret = EVP_PKEY_decrypt_init(this->ctx);
        if (ret <= 0)
        {
            throw runtime_error("can not init decrypt");
        }

        // 获取明文大小
        size_t out_size{};
        const auto* data_ptr = reinterpret_cast<const unsigned char*>(data.c_str());
        ret = EVP_PKEY_decrypt(ctx, nullptr, &out_size, data_ptr, data.length());
        if (ret <= 0)
        {
            throw runtime_error("can not get decrypt size");
        }

        // 解密
        string output(out_size + 1, '\0');
        auto out_ptr = reinterpret_cast<unsigned char*>(output.data());
        ret = EVP_PKEY_decrypt(ctx, out_ptr, &out_size, data_ptr, data.length());
        if (ret <= 0)
        {
            // 获取错误信息
            unsigned long err_code = 0;
            while ((err_code = ERR_get_error()))
            {
                char* err_str = ERR_error_string(err_code, nullptr);
                fprintf(stderr, "Decryption error: %s\n", err_str);
            }
            throw runtime_error("can not decrypt data");
        }
        return output.c_str();
    }
}
